#####################
# CAR_RACING V1 ENV #
#####################

import random
import time
import numpy as np
import pygame
import math
import gym
from enum import Enum
from typing import Dict, Tuple, Sequence, List
from utils.utils import *
from copy import deepcopy

# ----------------- META_DATA FOR TEST-------------------#
# WIDTH, HEIGHT = 1536, 864
WIDTH, HEIGHT = 1920, 1440

# CAR_POS_X, CAR_POS_Y = 690, 740
ORIGIN_POS_X, ORIGIN_POS_Y = WIDTH/2, HEIGHT/2

BORDER_COLOR = (255, 255, 255, 255)  # Color To Crash on Hit

PATH_TO_VEH = '/root/code/workspace/src//MVS/Materials/plane.png'

# TODO: 时间调整
FPS = 30 # 30倍放大
# TODO：比例缩放
wingspan = 1.363 # m # 翼展
LIM_SIZE = wingspan * FPS * 1.5
LIM_VEL = [15, 45]   # m/s  # 30倍放大
LIM_ANGVEL = [-15, 15]  # deg / s   # 30倍放大
LIM_ACC = [-1, 1]   # vel & ang


SWARM_NUM = 10
LEADER_ID = 1

class ColorTable(Enum):
    pass

class SurroundType(Enum):
    target = 0
    friend = 1
    obstacle = 2
    wall = 3
    
class renderMode(Enum):
    fixed = 0
    moved = 1

# ----------------- BASIC FUNCTIONS -------------------#
class myPoint:
    def __init__(self, x, y):
        self.x = x
        self.y = y

class myLine:
    def __init__(self, pt1, pt2):
        self.pt1 = myPoint(pt1.x, pt1.y)
        self.pt2 = myPoint(pt2.x, pt2.y) 

# class swarmInfo:
#     def __init__(self, my_id, leader_id, formation_type):
#         self.my_id = my_id
#         self.leader_id = leader_id
#         self.formation_type = formation_type

def distance(pt1, pt2):
    """ 欧式距离
    """
    return ((pt1.x - pt2.x) ** 2 + (pt1.y - pt2.y) ** 2) ** 0.5

def rotate(origin, point, angle, input_format='DEG'):
    """ LINE_ROTATOR
    """
    if not input_format == 'RAD':
        angle = math.radians(angle)
    qx = origin.x + math.cos(angle) * (point.x - origin.x) + math.sin(angle) * (point.y - origin.y)
    qy = origin.y - math.sin(angle) * (point.x - origin.x) + math.cos(angle) * (point.y - origin.y)
    q = myPoint(qx, qy)
    return q

def accelerate(x, d_x, bound):
    """ 加速度限幅
    """
    x_pass = x + d_x
    if x_pass > bound[1]:
        x_pass = bound[1]
    if x_pass < bound[0]:
        x_pass = bound[0]
    return x_pass

# ----------------- DYNAMIC ENV: TARGET CLASS -------------------#
# 建立动态环境：动态任务点
class Target(object):
    """ 2d 任务点（静态+动态）
    """
    def __init__(self, 
                 target_id : int = 1,
                 valid_pos : list = None,# 对象真实位置（中心点），
                 spawn_pos : list = [0., 0.],
                 designed_path : Sequence[Tuple[float, float]] = None,    # 预期路径, 默认为None(静态)
                 designed_color : Tuple[int, int, int] = (52,179,52),
                 max_hist_box: int = 25):
        
        self.target_id =target_id
        self.spawn_pos = spawn_pos
        self.valid_pos = valid_pos
        self.designed_path = designed_path  # 动态路径
        
        self.object_color = designed_color
        self.tar_hist = np.array([valid_pos])  # 目标点历史
        self.tar_traj = np.array([0., 0.])
        self.tar_hist_color = (84, 255, 159)
        self.max_hist_box = max_hist_box
        
    def update(self, valid_pos : list):
        self.valid_pos = deepcopy(valid_pos)
        # print(self.tar_hist[-1])
        # print(self.valid_pos)
        # if self.valid_pos != self.tar_hist[-1]:
        #     self.tar_hist = np.append(self.tar_hist, [self.valid_pos], axis=0)
        #     if self.tar_hist.shape[0] > self.max_hist_box:
        #         self.tar_hist = np.delete(self.tar_hist, obj=0, axis=0)
        #     self.tar_traj = self.tar_hist.copy()
        #     self.tar_traj[:, 0] = self.tar_hist[:, 0] - self.valid_pos[0]
        #     self.tar_traj[:, 1] = self.tar_hist[:, 1] - self.valid_pos[1]
        
    def spawn_to_pos(self, spawn_pos : list):
        self.spawn_pos = spawn_pos
        
    def render(self, screen):
        pygame.draw.circle(screen, self.object_color, self.spawn_pos, 10)
        # 绘制历史目标点
        # spawn_poses = []
        # for i in range(self.tar_traj.shape[0]):
        #     temp_pos = convert_coordinates(self.tar_traj[i], self.spawn_pos, is_flip=True)
        #     spawn_poses.append(temp_pos.tolist()[0])
        # spawn_poses.reverse()
        # for spawn_pos in spawn_poses:
        #     print("spawn_pos: ", spawn_pos)
        #     pygame.draw.circle(screen, self.object_color, spawn_pos, 8)
        
class Obstacle(object):
    """ 2d 障碍物（静态+动态）
    """
    def __init__(self, 
                 obstacle_id : int = 1,
                 valid_pos : list = None,# 对象真实位置（中心点），
                 spawn_pos : list = [0., 0.],
                 designed_path : Sequence[Tuple[float, float]] = None,    # 预期路径, 默认为None(静态)
                 designed_color : Tuple[int, int, int] = (52,179,52)):
        
        self.obstacle_id =obstacle_id
        self.spawn_pos = spawn_pos
        self.valid_pos = valid_pos
        self.designed_path = designed_path  # 动态路径
        
        self.object_color = designed_color
        
    def spawn_to_pos(self, spawn_pos : list):
        self.spawn_pos = spawn_pos
        
    def render(self, screen):
        pygame.draw.circle(screen, self.object_color, (self.spawn_pos[0], self.spawn_pos[1]), 10)        


# ----------------- DYNAMIC ENV: VEHICLE CLASS -------------------#
# 建立动态环境：载具类
class Vehicle(object):
    def __init__(self,
                 default_init : Tuple[int, float, float, float, float] = (1, 0.,0.,0.,0.),  # valid (id, pos_x, pos_y, vel, att)
                 default_spawn : Tuple[float, float] = (ORIGIN_POS_X, ORIGIN_POS_Y),  # render pos
                 image: str = PATH_TO_VEH,
                 size_bound: int = LIM_SIZE,
                 vel_bound: int = LIM_VEL,
                 angvel_bound: int = LIM_ANGVEL,
                 render_mode : renderMode = renderMode.fixed,
                 leader_mode: bool =False):
        """ 2d Vehicle 质点模型

        Args:
            default_init (Tuple[int, float, float, float, float], optional) : Vehicle 初始值. Defaults to (1, 0.,0.,0.,0.).
            default_spawn (Tuple[float, float])                             : Vehicle 渲染位置。Defaults to (ORIGIN_POS_X, ORIGIN_POS_Y).
            render_mode (renderMode, optional)                              : Vehicle 渲染模式(固连画面或移动). Defaults to renderMode.fixed.
        """
        # super().__init__()
        # Coordinator: pixel
        # Load Sprite and Rotate
        self.sprite = pygame.image.load(image).convert()  # Convert Speeds Up A Lot
        self.default_spawn_pos = myPoint(default_spawn[0], default_spawn[1])  # 画布展示位置
        self.spawn_pos = self.default_spawn_pos
        self.vehicle_id = default_init[0]
        
        # Other Vehicle Parameters
        self.render_mode = render_mode
        self.leader_mode = leader_mode
        
        # 改变载具图像大小
        self.veh_w = size_bound
        self.veh_h = int(self.veh_w / self.sprite.get_width() * self.sprite.get_height())
        self.sprite = pygame.transform.scale(self.sprite, (self.veh_w, self.veh_h))
        
        # 载具图像与姿态固连
        self.rotated_sprite = self.sprite
        self.rotated_sprite.set_colorkey((0, 0, 0))
        self.rect = self.rotated_sprite.get_rect() # .move(self.position.x, self.position.y)  # move to center (x,y)
        self.rect.center = self.spawn_pos.x, self.spawn_pos.y
        # self.center = myPoint(self.rect.centerx, self.rect.centery)  # Calculate Center

        # kinetic-LOCAL_NED
        # 真实位置
        self.default_position = myPoint(default_init[1], default_init[2]) 
        self.position = self.default_position    # 重置真实位置

        # 真实速度
        self.vel = default_init[4]
        self.default_vel = self.vel # 重置速度标量
        self.vel_bound = vel_bound
        self.dvel = 0.
        self.vecX = 0.
        self.vecY = 0.
        # 真实角度
        self.yaw = default_init[3]    # [-180, 180]
        self.default_yaw = self.yaw # 重置角度标量
        self.angvel = 0.
        self.angvel_bound = angvel_bound
        self.dangvel = 0.
        
        # Vehicle feedback
        self.isAlive = True  # Boolean To Check If Vehicle is Crashed
        self.time = 0  # Time Passed
        self.dist_driven = 0.  # Distance Driven
        self.history = np.array([[self.position.x, self.position.y]])
        self.trajectory = self.history
        self.vel_hist = np.array([self.vel])
        self.dis_hist = np.array([[0., 0.]])
        self.act_hist = [0., 0.]
    
    def rotate_image(self, image, angle):
        # Rotate The Rectangle
        rectangle = image.get_rect()
        rotated_image = pygame.transform.rotate(image,  - angle)
        rotated_rectangle = rectangle.copy()
        rotated_rectangle.center = rotated_image.get_rect().center
        return rotated_image
    
    def spawn_to_pos(self, spawn_init : Tuple[float, float]):
        self.spawn_pos = myPoint(spawn_init[0], spawn_init[1])
        
    def reset(self):

        self.spawn_pos = deepcopy(self.default_spawn_pos)   # 需要deepcopy防止指针内容被篡改
        
        self.vel = deepcopy(self.default_vel)
        self.dvel = 0.
        self.yaw = deepcopy(self.default_yaw)
        self.angvel = 0.
        self.dangvel = 0.
        
        self.position = deepcopy(self.default_position)     # 需要deepcopy防止指针内容被篡改
        
        self.isAlive = True
        self.time = 0  # Time Passed
        self.dist_driven = 0.  # Distance Driven
        self.history = np.array([[self.position.x, self.position.y]])
        self.trajectory = deepcopy(self.history)
        
        return [self.position.x, self.position.y, self.vel, self.yaw]
        
    # 运动
    def move(self, action):
        # 质点运动学模型
        print("Action: ",action)
        self.dvel, self.dangvel = action    # 
        if np.isnan(self.dvel) or np.isnan(self.angvel):
            self.dvel, self.angvel = self.act_hist[0], self.act_hist[1]
        else:
            self.act_hist[0] = deepcopy(self.dvel)
            self.act_hist[1] = deepcopy(self.angvel)
        self.vel = accelerate(self.vel, self.dvel, self.vel_bound)
        # self.angvel = accelerate(self.angvel, self.dangvel, self.angvel_bound)  # rad
        self.yaw = count_deg(self.yaw + math.degrees(self.dangvel))
        
        self.vecX = self.vel * math.cos(math.radians(self.yaw))
        self.vecY = self.vel * math.sin(math.radians(self.yaw))
        self.position.x += self.vecX
        self.position.y += self.vecY
        # 单机信息更新 state_raw
        return [self.position.x, self.position.y, self.vel, math.radians(self.yaw)]
    

    # 反馈更新
    def update(self, state_new):   # 接收邻居信息
        # Increase Distance and Time
        self.dist_driven += self.vel
        self.time += 1
        if self.time % 1 == 0:
            self.vel_hist = np.append(self.vel_hist, [self.vel], axis=0)
            self.history = np.append(self.history, [[self.position.x, self.position.y]], axis=0)
            # print([state_new[:2]])
            self.dis_hist = np.append(self.dis_hist, [state_new[:2]], axis=0)
            
            if self.history.shape[0] >= 50:
                self.history = np.delete(self.history, obj=0, axis=0)
        self.trajectory = self.history.copy()
        self.trajectory[:, 0] = self.history[:, 0] - self.position.x
        self.trajectory[:, 1] = self.history[:, 1] - self.position.y
        reward = 0.
        # TODO:更新存活状况，计算奖励 calc Reward_Solo （计算图）
        if not self.leader_mode:
            for id, val in state_new[4].items():
                # if np.linalg.norm(val[:2]) < 0.5 * FPS * wingspan  or np.linalg.norm(state_new[:2]) > 10. * FPS * wingspan:#* FPS * wingspan #(id == 1 and np.linalg.norm(val[:2]) > 20. * FPS * wingspan):   # 出界
                if np.linalg.norm(val[:2]) < 0.5 or np.linalg.norm(state_new[:2]) > 10.:
                    self.isAlive = False
                    reward_orig = -10.   # APF 10 或者离主机太远
                else:
                    reward_orig = 0.

            forma_r = reward_adaptor(state_new, FPS, wingspan, (self.dis_hist, self.vel_hist))
            # print("forma_r: ", forma_r)
            reward = reward_orig + forma_r

        return reward
        
    
    def trajectoryFilterDraw(self, screen):
        traj_offset = self.time % 4
        del_idx = []
        screen_rect = screen.get_rect()
        # if self.trajectory.shape[0] >=2:
        #     pygame.draw.lines(screen, (255,114,86), False, self.trajectory.tolist(),5)
        spawn_poses = []
        for i in range(self.trajectory.shape[0]):
            # print(self.vehicle.trajectory[i])
            # if self.time % 2 == 0 :
            #     pass
            temp_pos = convert_coordinates(self.trajectory[i],[self.spawn_pos.x, self.spawn_pos.y], is_flip=True)
            spawn_poses.append(temp_pos.tolist()[0])
            # print(spawn_pos[0,0], spawn_pos[0,1])
            # print(spawn_pos)
        if self.trajectory.shape[0] >=2:
            if self.vehicle_id == 1:
                pygame.draw.lines(screen, (255,114,86), False, spawn_poses,5)
            else:
                pygame.draw.lines(screen, (0,154,205), False, spawn_poses,5)
        spawn_poses.reverse()
        for i, spawn_pos in enumerate(spawn_poses):
            # print(spawn_pos)    
            if screen_rect.left < spawn_pos[0] < screen_rect.right and screen_rect.top < spawn_pos[1] < screen_rect.bottom:
                if i % 5 == traj_offset:
                    radius = 12
                else:
                    radius = 6
                if self.vehicle_id == 1:
                    pygame.draw.circle(screen, (255,255,255), spawn_pos, radius)
                else:
                    pygame.draw.circle(screen, (255,255,255), spawn_pos, radius)

            else:
                del_idx.append(i)
        # for idx in del_idx:
        #     self.vehicle.trajectory = np.delete(self.vehicle.trajectory, obj=idx, axis=0)
    
    def render(self, screen):
        # 绘制历史轨迹
        self.trajectoryFilterDraw(screen)
        # 绘制速度方向（箭头）
        arrow_color = (230,200,101)
        arrow_len = 5
        arrow_wid = 5
        start_pos_tuple = (self.spawn_pos.x, self.spawn_pos.y)
        # print(arrow_start_pos)
        trunk_end = convert_coordinates([self.vecX * 3, self.vecY * 3], start_pos_tuple, is_flip=True)
        arrow_end = convert_coordinates([self.vecX * 3 + arrow_len * math.cos(math.radians(self.yaw)), self.vecY * 3 + arrow_len * math.sin(math.radians(self.yaw))], start_pos_tuple, is_flip=True)
        arrow_left = convert_coordinates([self.vecX * 3 - arrow_len * math.cos(math.radians(self.yaw)) + arrow_wid * math.sin(math.radians(self.yaw)), self.vecY * 3 - arrow_len * math.sin(math.radians(self.yaw)) - arrow_wid * math.cos(math.radians(self.yaw))], start_pos_tuple, is_flip=True)
        arrow_right = convert_coordinates([self.vecX * 3 - arrow_len * math.cos(math.radians(self.yaw)) - arrow_wid * math.sin(math.radians(self.yaw)), self.vecY * 3 - arrow_len * math.sin(math.radians(self.yaw)) + arrow_wid * math.cos(math.radians(self.yaw))], start_pos_tuple, is_flip=True)
        end_pos_tuple = (trunk_end[0,0], trunk_end[0,1])
        arr_end_pos = (arrow_end[0,0], arrow_end[0,1])
        arr_left_pos = (arrow_left[0,0], arrow_left[0,1])
        arr_right_pos = (arrow_right[0,0], arrow_right[0,1])
        
        # print(arrow_end_pos)
        if not end_pos_tuple == start_pos_tuple:
            pygame.draw.line(screen, arrow_color, start_pos_tuple, end_pos_tuple, 2)
            pygame.draw.polygon(screen, arrow_color, (arr_end_pos, arr_left_pos, arr_right_pos))
        # 绘制飞行器图标
        # Get Rotated Sprite And Move Into The Right X-Direction
        self.rotated_sprite = self.rotate_image(self.sprite, self.yaw)
        self.rect = self.rotated_sprite.get_rect()  # Replace old rect with new rect.
        self.rect.center = self.spawn_pos.x, self.spawn_pos.y
        screen.blit(self.rotated_sprite, self.rect)
        
        
# ----------------- Establish ENV: Map CLASS -------------------#
# 建立环境：单机感知范围观测图
class SoloObservationMap(object):
    def __init__(self, vehicle : Vehicle, 
                 default_target_info : Tuple[float, float] = None, 
                 default_surround_info : Sequence[Tuple[float, float]] = None, 
                 default_spawn : Tuple[float, float, float, float] = (ORIGIN_POS_X, ORIGIN_POS_Y, 0., 0.),
                 state_dim = 4, action_dim = 2, action_bound = [-1, 1],
                 fps=FPS, width=WIDTH, height=HEIGHT, default_render_fixed_id: int = 1):
        
        self.vehicle = None
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.action_bound = action_bound
        
        self.default_pos = myPoint(default_spawn[0], default_spawn[1])
        self.default_vel = default_spawn[2]
        self.default_att = default_spawn[3]
        
        self.default_target_info = default_target_info  # 数据类型 TargetInfo
        self.default_surround_info = default_surround_info
        
        pygame.init()
        # self.font = pygame.font.Font(pygame.font.get_default_font(), 36)
        self.generation_font = pygame.font.SysFont("Arial", 20)
        self.alive_font = pygame.font.SysFont("Arial", 10)
        self.fps = fps
        self.width = width
        self.height = height
        
        # 生成画布及范围
        self.screen = pygame.display.set_mode((self.width, self.height))
        pygame.display.set_caption("Solo Observation Map")
        self.screen.fill((0, 0, 0))
        self.screen_rect = self.screen.get_rect()

        self.action_space = None
        self.observation_space = None
        self.game_reward = 0
        self.history = []
        self.time = 0
        
        self.vehicle = vehicle
        # self.fixed_id = default_render_fixed_id
        self.target = Target(object_id=0, 
                                     surround_type=SurroundType.target, 
                                     spawn_pos=self.default_target_info)
        
        
        self.reset()

    def reset(self):
        print("RESET!")
        # pygame.init()
        pygame.event.get()
        # self.screen.fill((90, 90, 90))
        # 载具重置及可视化
        for i, _ in enumerate(self.vehicles):
            self.vehicles[i].reset()
        
        self.game_reward = 0
        self.time = 0

    def step(self, action):
        self.time += 1
        done = False
        self.vehicle.move(action)
        state_new = [self.vehicle.position.x, self.vehicle.position.y, self.vehicle.vel, self.vehicle.yaw]
        if self.vehicle.isAlive and self.time <=200:
            self.vehicle.update()
        else:
            done = True
            state_new = None

        return state_new, self.game_reward, done

    def render(self):
        pygame.event.get()
        pygame.time.delay(10)
        self.clock = pygame.time.Clock()
        self.screen.fill((255, 255, 255))
        # Draw Map And All Vehicles That Are Alive

        if self.vehicle.isAlive:
            self.target.render(self.screen)
            self.trajectoryFilterDraw()
            self.vehicle.render(self.screen)
            
            
        # Display Info
        # text_surface = self.generation_font.render(f'Speed {self.vehicle.vel * -1}', True, pygame.Color('green'))
        # self.screen.blit(text_surface, dest=(400, 0))
        # text2 = self.generation_font.render(f'Reward {self.game_reward}', True, pygame.Color('red'))
        # self.screen.blit(text2, dest=(400, 50))

        self.clock.tick(self.fps)
        pygame.display.flip()

    def trajectoryFilterDraw(self):
        del_idx = []
        for i in range(self.vehicle.trajectory.shape[0]):
            # print(self.vehicle.trajectory[i])
            spawn_pos = convert_coordinates(self.vehicle.trajectory[i],[self.default_pos.x, self.default_pos.y], is_flip=True)
            # print(spawn_pos[0,0], spawn_pos[0,1])
            if self.screen_rect.left < spawn_pos[0,0] < self.screen_rect.right and self.screen_rect.top < spawn_pos[0,1] < self.screen_rect.bottom:
                pygame.draw.circle(self.screen, (138,116,46), (spawn_pos[0,0], spawn_pos[0,1]), 10)
            else:
                del_idx.append(i)
        # for idx in del_idx:
        #     self.vehicle.trajectory = np.delete(self.vehicle.trajectory, obj=idx, axis=0)
    
    def close(self):
        pygame.quit()

    def sample_actions(self):
        return np.clip([random.uniform(-1, 1), random.uniform(-1, 1)], *self.action_bound)


# ----------------- Establish ENV: Map CLASS -------------------#
# 建立环境：全局感知范围观测图
class GeneralObservationMap(object):
    def __init__(self, screen_init, vehicles : Sequence[Vehicle],
                 formation_set : dict = None, 
                 default_surround_info : Sequence[Tuple[float, float]] = None, 
                 default_spawn : Tuple[float, float, float, float] = (ORIGIN_POS_X, ORIGIN_POS_Y, 0., 0.),
                 state_dim = 4, action_dim = 2, action_bound = [-1, 1],
                 fps=FPS, width=WIDTH, height=HEIGHT, default_render_fixed_id: int = 1):
        
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.action_bound = action_bound
        
        self.default_pos = myPoint(default_spawn[0], default_spawn[1])
        self.default_vel = default_spawn[2]
        self.default_att = default_spawn[3]
        
        # self.default_target_info = default_target_info  # 数据类型 TargetInfo
        self.default_surround_info = default_surround_info
        
        pygame.init()
        # self.font = pygame.font.Font(pygame.font.get_default_font(), 36)
        self.generation_font = pygame.font.SysFont("Arial", 20)
        self.alive_font = pygame.font.SysFont("Arial", 10)
        self.fps = fps
        self.width = width
        self.height = height
        
        # 生成画布及范围
        # self.screen = pygame.display.set_mode((self.width, self.height))
        # pygame.display.set_caption("General Observation Map")
        # self.screen.fill((0, 0, 0))
        self.screen = screen_init
        self.screen_rect = self.screen.get_rect()

        self.action_space = None
        self.observation_space = None
        self.game_reward = {}
        self.history = []
        self.time = 0
        
        self.vehicles = vehicles
        self.formation_set = formation_set
        self.leader_id = self.formation_set.leader_id
        self.default_fixed_id = default_render_fixed_id
        
        self.formaDST = DST(None)   # 编队目标点的DST管理结构
        self.swarmDST = DST(None)   # 集群节点的DST管理结构
        self.vehicle_nodes_dict = {}
        
        self.reset()

    def createDST(self, veh_params: list):
        # 创建DST管理结构
        edges_list, nodes_dict = generate_edges_list(veh_params)
        baseDST, f_nodes_dict, formation_sp = create_formation_baseDST(veh_params, self.formation_set)
        graph = create_graph(edges_list)
        validDST, mst_edges, mst_weight, new_nodes_dict = formation_prim_tree(graph, nodes_dict, baseDST)
        return baseDST, validDST, new_nodes_dict
    
    def rebuild_forma_offset_id(self, p: 'Node[T]') -> int:
        fr_offset = p.get_fromRoot_id()
        if self.formation_set.formation_type == FormationType.wingLike.value:
            offset_id = fr_offset * 2 + (p.get_value()[0] % 2)
        else:
            offset_id = fr_offset
        return offset_id
    
    def reset(self):
        print("RESET!")
        # pygame.init()
        pygame.event.get()
        # self.screen.fill((90, 90, 90))
        # 载具重置及可视化
        self.targets, state_raw, targets_pos, swarm_info = {}, {}, {}, []
        for id, veh in self.vehicles.items():
            x_offset, y_offset = 0., 0.
            state_raw[id] = self.vehicles[id].reset()
            swarm_info.append([int(id), state_raw[id][0], state_raw[id][1]])
            
            self.game_reward[id] = 0.
        
        # 建立集群成员管理结构
        self.formaDST, self.swarmDST, self.vehicle_nodes_dict = self.createDST(swarm_info)

        # 重置目标队形 
        for id, veh in self.vehicles.items():
            # x_offset, y_offset = calc_formation_offset(self.formation_set, self.vehicle_nodes_dict[id][1].get_value()[0], 0.)
            x_offset, y_offset = calc_formation_offset(self.formation_set, self.rebuild_forma_offset_id(self.vehicle_nodes_dict[id][1]), 0.)
            self.targets[id] = Target(id, (self.vehicles[self.leader_id].default_position.x + x_offset, self.vehicles[self.leader_id].default_position.y + y_offset))
            targets_pos[id] = self.targets[id].valid_pos
        
        state_reset = state_format(state_raw, targets_pos)  # 计算与周围成员的距离差
        self.time = 0
        
        # 创建log，存位置、速度、角度(画图用，RL不用)
        
        
        return state_reset

    def step(self, action : dict):
        removed_id = -1
        self.time += 1
        done, terminate, state_raw, state_new, targets_pos, reward = {}, {}, {}, {}, {}, {}
        isAliveList = [self.vehicles[id].isAlive for id, _ in self.vehicles.items()]
        
        if len(action) == 0:    # RESET后，没有接到初始动作信息 # 不要用 == {} 来判断，有时候无效 或所有僚机均未生还
            for id, veh in self.vehicles.items():
                done[id], terminate[id] = False, False
                state_new = self.reset()
                reward[id] = 0. # 单机初始reward
        else:
            for id, veh in self.vehicles.items():
                if self.vehicles[id].isAlive:
                    x_offset, y_offset = 0., 0.
                    done[id], terminate[id] = False, False
                    state_raw[id] = self.vehicles[id].move(action[id][0:2]) # [self.vehicles[id].position.x, self.vehicles[id].position.y, self.vehicles[id].vel, self.vehicles[id].yaw]

                    # 更新目标位置 长机计算目标点位置
                    # print(id)
                    # print(self.vehicle_nodes_dict[id])
                    # x_offset, y_offset = calc_formation_offset(self.formation_set, self.vehicle_nodes_dict[id][1].get_value()[0], state_raw[self.leader_id][3])    # 计算对应编队目标点
                    x_offset, y_offset = calc_formation_offset(self.formation_set, self.rebuild_forma_offset_id(self.vehicle_nodes_dict[id][1]), state_raw[self.leader_id][3])    # 计算对应编队目标点

                    self.targets[id].update([self.vehicles[self.leader_id].position.x + x_offset, self.vehicles[self.leader_id].position.y + y_offset])
                    # self.targets[id].valid_pos = [self.vehicles[self.leader_id].position.x + x_offset, self.vehicles[self.leader_id].position.y + y_offset]
                                    
                    # 分布式目标点设置
                    # x_offset, y_offset = calc_formation_offset(self.formation_set, id % 2 + 2, state_raw[self.leader_id][3])    # 计算对应编队目标点
                    # self.targets[id].valid_pos = [self.vehicles[self.leader_id if (id-2) <=0 else id-2].position.x + x_offset, self.vehicles[self.leader_id if (id-2) <=0 else id-2].position.y + y_offset]
                    targets_pos[id] = deepcopy(self.targets[id].valid_pos)
            # 计算与周围成员的距离差
            state_new = state_format(state_raw, targets_pos)
            # 更新奖励等信息
            for id, veh in self.vehicles.items():
                if not self.vehicles[id].isAlive:
                    removed_id = id
                    
                    # 在树状结构中管理（移除）节点
                    if id in self.vehicle_nodes_dict.keys():
                        # print(self.vehicle_nodes_dict.items())
                        # print(self.vehicle_nodes_dict[id][0].get_value())
                        # print(self.vehicle_nodes_dict[id][1].get_value())
                        # print(self.swarmDST.get_subtree(self.swarmDST.get_root()))
                        # print(self.formaDST.get_subtree(self.formaDST.get_root()))
                        self.swarmDST.remove_node(self.vehicle_nodes_dict[id][0])
                        self.formaDST.remove_node(self.vehicle_nodes_dict[id][1])
                        
                        self.vehicle_nodes_dict.pop(id)
                        
                    
                    done[id], terminate[id] = True, True
                    state_new[id] = [None for i in range(5)]
                    reward[id] = -10.# self.game_reward[id] # 单机总reward
                
                if self.vehicles[id].isAlive and self.time <= 3000:  # 100s
                    reward[id] = self.vehicles[id].update(state_new[id]) # and isAlive
                    # print(reward[id])
                    
                    # self.game_reward[id] += reward[id]
                elif self.time > 3000:
                    terminate[id] = True
                    reward[id] = 0.
                    
            for id, veh in self.vehicles.items():
                print(id, (done[id] or terminate[id]), reward[id])
                if id != removed_id and removed_id != -1 and (done[id] or terminate[id]) != True:  # state_new[id][4] != None
                    # print(state_new[id][4])
                    if removed_id in state_new[id][4].keys():
                        state_new[id][4].pop(removed_id)
                    
        if isAliveList.count(True) < 2:
            for id, veh in self.vehicles.items():
                done[id], terminate[id] = False, True

        return state_new, reward, done, terminate

    def render(self):
        pygame.event.get()
        pygame.time.delay(10)
        self.clock = pygame.time.Clock()
        self.screen.fill((255, 255, 255))
        # Draw Map And All Vehicles That Are Alive
        
        for id, veh in self.vehicles.items():
            
            if self.vehicles[id].isAlive:
                # vehicle spawn pos
                spawn_pos = convert_coordinates((veh.position.x - self.vehicles[self.default_fixed_id].position.x, veh.position.y - self.vehicles[self.default_fixed_id].position.y),(960,720), is_flip=True)
                self.vehicle_nodes_dict[id][0].set_value((id, spawn_pos[0,0], spawn_pos[0,1]))
                veh.spawn_to_pos((spawn_pos[0,0], spawn_pos[0,1]))
                # targets initialize
                if id != self.default_fixed_id:    
                    spawn_posT = convert_coordinates((self.targets[id].valid_pos[0] - self.vehicles[self.default_fixed_id].position.x, self.targets[id].valid_pos[1] - self.vehicles[self.default_fixed_id].position.y),(960,720), is_flip=True)
                    self.vehicle_nodes_dict[id][1].set_value((self.vehicle_nodes_dict[id][1].get_value()[0], spawn_posT[0,0], spawn_posT[0,1]))
                    self.targets[id].spawn_to_pos([spawn_posT[0,0], spawn_posT[0,1]])
                
                # 渲染机体、目标点
                self.vehicles[id].render(self.screen)
                self.targets[id].render(self.screen)
        
        # 绘制DST结构
        if self.swarmDST.get_size() > 1:
            root = self.swarmDST.get_root()
            visited = [root]
            while len(visited) < self.swarmDST.get_size():
                for u_node in visited:
                    if u_node.is_internal():
                        for v_node in u_node.get_children():
                            if v_node not in visited:
                                pygame.draw.line(self.screen, (135,206,250), u_node.get_value()[1:], v_node.get_value()[1:], 4)
                                visited.append(v_node)
            
        for idx, node in self.vehicle_nodes_dict.items():
            if self.vehicles[idx].isAlive:
                pygame.draw.circle(self.screen,(255,127,80) if idx == 1 else (100,149,237),node[0].get_value()[1:],10)
                pygame.draw.circle(self.screen,(255,127,80) if idx == 1 else (135,206,250),node[0].get_value()[1:],10,2)
                font = pygame.font.Font(None, 35)
                text_surface = font.render(str(idx), True, (0,0,0))
                text_width, text_height = text_surface.get_size()
                text_x = node[0].get_value()[1] - (text_width / 2)
                text_y = node[0].get_value()[2] - (text_height / 2)
                self.screen.blit(text_surface, (text_x, text_y))
            
        # TODO: Display Info
        # text_surface = self.generation_font.render(f'Speed {self.vehicle.vel * -1}', True, pygame.Color('green'))
        # self.screen.blit(text_surface, dest=(400, 0))
        # text2 = self.generation_font.render(f'Reward {self.game_reward}', True, pygame.Color('red'))
        # self.screen.blit(text2, dest=(400, 50))

        self.clock.tick(self.fps)
        pygame.display.flip()
    
    def add_vehicle_node(self):
        # 中途加载具
        pass
    
    def close(self):
        pygame.quit()

    def sample_actions(self):
        actions = {}
        for id, veh in self.vehicles.items():
            actions[id] = np.clip([random.uniform(-1, 1), random.uniform(-1, 1)], *self.action_bound)
        return actions

# # 测试程序——随机运动
if __name__ == "__main__":
    pass
    
#     env = GeneralObservationMap(id=1, default_target_info=(ORIGIN_POS_X, ORIGIN_POS_Y * 0.25))
#     ep_r = 0
#     while True:
#         env.render()
#         a = env.sample_actions()
#         # a = [5.,0.1]
#         s, r, done = env.step(a)
#         ep_r += r
#         # print("State:",s)
#         if done:
#             s = env.reset()
#             print("Reward:", ep_r)
#             ep_r = 0
#         for event in pygame.event.get():
#             # print(1)
#             if event.type == pygame.QUIT:
#                 env.close()
